#!/usr/bin/env python


import roslib; roslib.load_manifest('knobs_rev')
import rospy
import message_filters
import image_geometry
import actionlib
import tf
import tf.transformations as tt

from smach import State, StateMachine
from sensor_msgs.msg import Image, CameraInfo, PointCloud2
from arm_navigation_msgs.msg import MoveArmGoal, MoveArmAction, JointConstraint
from pr2_controllers_msgs.msg import Pr2GripperCommandGoal, Pr2GripperCommandAction
from std_msgs.msg import Float32, Header

from knobs_rev.msg import PoseArmAction, PoseArmGoal
from knobs_rev.srv import PlaneCloud, PlaneCloudRequest, PlaneCloudResponse

import cv
from cv_bridge import CvBridge, CvBridgeError
from PySide.QtCore import * 
from PySide.QtGui import * 

import sys
from time import sleep
import threading

#************************** HELPERS ************************************
# Convert cv image to qimage... Borrowed from somewhere online, which you can find if you google the class name
class IplQImage(QImage):
   def __init__(self, iplimage):
      rgba = cv.CreateMat(iplimage.height, iplimage.width, cv.CV_8UC4)
      cv.Set(rgba, (1,2,3,4))

      cv.MixChannels([iplimage], [rgba], [ (0,0), (1,1), (2,2), (3,3) ] ) # This encoding works because I explicitly convert the imgmsg to cv in rgba
      self.__imagedata = rgba.tostring()
      super(IplQImage,self).__init__(self.__imagedata, iplimage.width, iplimage.height, QImage.Format_RGB32)

class QSignal(QObject):
   signal = Signal()

class MouseSignal(QObject):
   signal = Signal(int, int, int, int)

class MyLabel(QLabel):
   def __init__(self):
      super(MyLabel, self).__init__()
      self.mousePress = False
      self.start = None
      self.end = None
      self.done = MouseSignal()

   def mousePressEvent(self, event):
      if self.mousePress:
         (ex, ey) = (event.x(), event.y())
         (sx, sy) = self.start
         self.mousePress = False
         self.done.signal.emit(sx, sy, ex, ey)
      else:
         self.start = (event.x(), event.y())
         self.mousePress = True
      event.accept()

   def mouseMoveEvent(self, event):
      event.accept()


#************************** STATES ************************************
class ClearArm(State):
   def __init__(self):
      State.__init__(self, outcomes=['succeeded', 'abort'])
      self.right_arm_client = actionlib.SimpleActionClient('pose_arm', PoseArmAction)
      rospy.loginfo("Waiting on arm server")
      self.right_arm_client.wait_for_server()
      rospy.loginfo("Arm server found")
        

   def execute(self, userdata):
      rospy.loginfo("Clear arm state")  
      if rospy.is_shutdown():
         return 'abort'

      goal = PoseArmGoal()
      goal.pose.header.frame_id = 'torso_lift_link'
      goal.pose.pose.position.x = 0.5
      goal.pose.pose.position.y = -0.3
      goal.pose.pose.position.z = -0.1
      goal.pose.pose.orientation.x = 0.0
      goal.pose.pose.orientation.y = 0.0
      goal.pose.pose.orientation.z = 0.0
      goal.pose.pose.orientation.w = 1.0

      self.right_arm_client.send_goal(goal)
      rospy.loginfo("Waiting for arms to clear")
      self.right_arm_client.wait_for_result()
      return 'succeeded'


class Click(State):
    def __init__(self, group_name, buttons):
        State.__init__(self, outcomes=['succeeded', 'failed', 'abort'])
        self.layout = QVBoxLayout()

        self.button_list = [QPushButton(name) for name in buttons]
        self.return_list = [value for value in buttons.values()]
        map(lambda x: x[0].clicked.connect(lambda: self.callback(x[1])), zip(self.button_list, self.return_list))
        map(self.layout.addWidget, self.button_list)

        self.group = QGroupBox(group_name)
        self.group.setLayout(self.layout)

    def callback(self, value):
        self.return_value = value
        self.group.hide()
        app.quit()

    def execute(self, userdata):
        rospy.loginfo("Click state")
        self.return_value = None

        self.group.show()
        app.exec_()

        return self.return_value

class NoDisplay(State):
   def __init__(self, image_feed):
      State.__init__(self, outcomes=['succeeded', 'failed', 'abort'], output_keys=['low', 'high'])

   def execute(self, userdata):
      rospy.loginfo("Skipping display")
      userdata.low = (200,200)
      userdata.high = (400, 400)
      return 'succeeded'


class Display(State):
   def __init__(self, image_feed):
      State.__init__(self, outcomes=['succeeded', 'failed', 'abort'],
                           output_keys=['low', 'high'])
      # Flags and global vars
      self.active = False
      self.userdata = None
      self.qimage = QImage()
      # OpenCv setup
      self.bridge = CvBridge()
      # Publishers and Subscribers
      self.image_sub = rospy.Subscriber(image_feed, Image, self.callback)
      # Pyside
      self.imageReady = QSignal()
      self.imageLock = threading.Lock()
      self.group = QGroupBox('display')
      self.label2 = MyLabel()

      self.button = QPushButton('done') 

      vbox = QVBoxLayout()
      vbox.addWidget(self.label2)
      vbox.addWidget(self.button)
      self.group.setLayout(vbox)

   def done(self, sx=None, sy=None, ex=None, ey=None):
      self.userdata = ((sx, sy), (ex, ey))
      if sx == None:
         self.userdata = (None, None)
      self.group.hide()
      app.quit()

   def execute(self, userdata):
      rospy.loginfo("Display state")
      self.imageReady.signal.connect(self.showImage)
      self.button.clicked.connect(self.done)
      self.label2.done.signal.connect(self.done)
      self.active = True
      self.group.show()
      app.exec_()

      rospy.loginfo("app executed")
      self.active = False
      (start, end) = self.userdata
      if start == None:
         return 'abort'
      (sx, sy) = start
      (ex, ey) = end
      userdata.high = (max(sx, ex), max(sy, ey))
      userdata.low = (min(sx, ex), min(sy, ey))
      return 'succeeded' 

   def showImage(self):
      #rospy.loginfo("showImage waiting to acquire lock")
      self.imageLock.acquire()
      qpixmap = QPixmap.fromImage(self.qimage)
      #rospy.loginfo("showImage releasing lock")
      self.label2.setPixmap(qpixmap)
      self.imageLock.release()

   def callback(self, data):
      if not self.active:
         return
      if rospy.is_shutdown():
         return
      try:
         cv_image = self.bridge.imgmsg_to_cv(data, "rgba8")
      except CvBridgeError, e:
         print e
         return
      #rospy.loginfo("callback waiting to acquire lock")
      self.imageLock.acquire()
      self.qimage = IplQImage(cv_image)
      #rospy.loginfo("callback releasing lock")
      self.imageLock.release()
      self.imageReady.signal.emit()


class CalculatePlane(State):
   def __init__(self, cloud_name, camera_info_name):
      State.__init__(self, outcomes=['succeeded', 'failed', 'abort'],
                           input_keys=['low', 'high'],
                           output_keys=['normal', 'point'])
      # Vars
      self.done = False
      self.cam_model = None
      self.cloud = None
      self.frame = None
      
      # Subs
      rospy.Subscriber(camera_info_name, CameraInfo, self.caminfo_callback)
      rospy.Subscriber(cloud_name, PointCloud2, self.cloud_callback)

      # Get a service proxy for the plane segmenter
      rospy.wait_for_service('plane_finder')
      self.plane_seg = rospy.ServiceProxy('plane_finder', PlaneCloud)

   def cloud_callback(self, cloud):
      if (self.cloud == None) and (not self.cam_model == None):
         self.cloud = cloud
         self.done = True

   def caminfo_callback(self, cam_info):
      if self.cam_model == None:
         self.frame = cam_info.header.frame_id
         self.cam_model = image_geometry.PinholeCameraModel()
         self.cam_model.fromCameraInfo(cam_info)
         

   def execute(self, userdata):
      rospy.loginfo("calculate plane callback")
      self.done = False
      self.cloud = None
      while not (self.done or rospy.is_shutdown()):
         rospy.sleep(0.01)
      if rospy.is_shutdown():
         return 'abort'

      (lx, ly) = userdata.low
      (hx, hy) = userdata.high

      # Request to plane finder
      request = PlaneCloudRequest()
      request.header = Header()
      request.header.stamp = rospy.Time(0)
      request.header.frame_id =  self.frame
      rospy.loginfo("self.cloud == None? {0}".format(self.cloud==None))
      request.cloud = self.cloud

      # Convert from image space to world space
      ul_ray = self.cam_model.projectPixelTo3dRay(userdata.low)
      ur_ray = self.cam_model.projectPixelTo3dRay((hx, ly))
      ll_ray = self.cam_model.projectPixelTo3dRay((lx, hy))
      lr_ray = self.cam_model.projectPixelTo3dRay(userdata.high)
      center_ray = self.cam_model.projectPixelTo3dRay(((lx+hx)/2.0 , (ly+hy)/2.0))

      # Fake min and max z # FIXME
      request.low_z = 0.5
      request.high_z = 1.5
      # Calculate the x, y vals based on closest z value 
            # Is this mathematically sound???
      request.low_x = min(ul_ray[0], ll_ray[0]) * request.low_z
      request.high_x = max(ur_ray[0], lr_ray[0]) * request.low_z
      request.low_y = min(ur_ray[1], ul_ray[1]) * request.low_z
      request.high_y = max(ur_ray[1], ul_ray[1]) * request.low_z

      # HACK # FIXME
      request.low_z = -100.0
      request.high_z = 100.0
      request.low_x = -100.0
      request.high_x = 100.0
      request.low_y = -100.0
      request.high_y = 100.0


      # Get plane from point cloud
      response = PlaneCloudResponse()
      response = self.plane_seg(request)
      
      # Intersection of plane and line (assumes that everyone is in the same frame)
      t = response.d / (response.a * center_ray[0] + response.b * center_ray[1] + response.c * center_ray[2])

      # Point at end of knob
      x = center_ray[0] * t
      y = center_ray[1] * t
      z = center_ray[2] * t

      # Passing to next state
      userdata.normal = (response.a, response.b, response.c)
      userdata.point = (x,y,z)
   
      return 'succeeded'


class PositionArm(State):
   def __init__(self):
      State.__init__(self, outcomes=['succeeded', 'failed', 'abort'],
                           input_keys=['normal', 'point'])
      self.right_arm_client = actionlib.SimpleActionClient('pose_arm', PoseArmAction)
      self.right_arm_client.wait_for_server()
      self.tf = tf.TransformListener()

   def execute(self, userdata):
      rospy.loginfo("position arm state")
      if rospy.is_shutdown():
         return 'abort'

      goal = PoseArmGoal() #FIXME -> Ignores userdata.point and userdata.normal
      goal.pose.header.frame_id = 'torso_lift_link'
      goal.pose.pose.position.x = 0.72
      goal.pose.pose.position.y = -0.18
      goal.pose.pose.position.z = 0.09

      rospy.loginfo("input to position arm: {0}, {1}".format(userdata.normal, userdata.point))
      
      #t = -0.2
      #goal.pose.header.frame_id = '/narrow_stereo_optical_frame'
      goal.pose.header.stamp = rospy.Time(0)
      #goal.pose.pose.position.x = userdata.point[0] #+ t * userdata.normal[0]
      #goal.pose.pose.position.y = userdata.point[1] #+ t * userdata.normal[1]
      #goal.pose.pose.position.z = userdata.point[2] #+ t * userdata.normal[2]

      #rospy.loginfo("Moving arm to: {0}".format(goal))
      self.right_arm_client.send_goal(goal)
      rospy.loginfo("waiting on arm result")
      self.right_arm_client.wait_for_result()
      return 'succeeded'


class FinalApproach(State):
   def __init__(self):
      State.__init__(self, outcomes = ['succeeded', 'failed', 'abort'])
      self.right_arm_client = actionlib.SimpleActionClient('/move_right_arm', MoveArmAction)
      rospy.loginfo("Waiting on move_right_arm server")
      self.right_arm_client.wait_for_server()

   def execute(self, userdata):
      rospy.loginfo("final approach state")
      if rospy.is_shutdown():
         return 'abort'
      goal = MoveArmGoal()
      goal.motion_plan_request.group_name = 'right_arm'
      goal.motion_plan_request.num_planning_attempts = 1
      goal.motion_plan_request.allowed_planning_time = rospy.Duration(5.0)
      goal.motion_plan_request.planner_id = ''
      goal.planner_service_name = 'ompl_planning/plan_kinematic_path'

      joint_angles = {'r_upper_arm_roll_joint':-1.5466936482606153, 
                        'r_shoulder_pan_joint':-0.24061559789458131, 
                        'r_shoulder_lift_joint':-0.13052543328701074, 
                        'r_forearm_roll_joint':-2.8959025637550835, 
                        'r_elbow_flex_joint':-0.67977712040984106, 
                        'r_wrist_flex_joint':-0.50094827692654909, 
                        'r_wrist_roll_joint':-1.7652659302215628} 

      for joint, angle in joint_angles.iteritems():
         j = JointConstraint()
         j.joint_name = joint
         j.position = angle
         j.tolerance_above = 0.01
         j.tolerance_below = 0.01
         j.weight = 1.0
         goal.motion_plan_request.goal_constraints.joint_constraints.append(j)

      for i in range(10):
         self.right_arm_client.send_goal(goal)
         self.right_arm_client.wait_for_result()

      return 'succeeded'


class Gripper(State):
   def __init__(self, close):
      State.__init__(self, outcomes=['succeeded', 'failed', 'abort'])
      self.close = close
      self.gripper_client = actionlib.SimpleActionClient('r_gripper_controller/gripper_action', Pr2GripperCommandAction)
      rospy.loginfo("Gripper state waiting on gripper client")
      self.gripper_client.wait_for_server()
      rospy.loginfo("Gripper state found gripper client")

   def execute(self, userdata):
      rospy.loginfo("gripper state")
      if rospy.is_shutdown():
         return 'abort'
      goal = Pr2GripperCommandGoal()
      if self.close:
         goal.command.position = 0.03
      else:
         goal.command.position = 0.08
   
      goal.command.max_effort = 50.0

      self.gripper_client.send_goal(goal)
      self.gripper_client.wait_for_result()
      
      return 'succeeded'

class TurnKnob(State):
    def __init__(self):
      State.__init__(self, outcomes=['succeeded', 'failed', 'abort'], input_keys=['rotation'])
      self.right_arm_client = actionlib.SimpleActionClient('pose_arm', PoseArmAction)
      rospy.loginfo("Waiting on arm server")
      #self.right_arm_client.wait_for_server()
      rospy.loginfo("Arm server found")
      #self.prev_angle = 0.0 # Not needed, arms default to set position

    def execute(self, userdata):
      if rospy.is_shutdown():
         return 'abort'
      if not userdata.rotation: # For some reason this actually can happen.
         return 'failed'
      rospy.loginfo("Turning knob to {0}".format(userdata.rotation))

      #change = userdata.rotation - self.prev_angle
      #self.prev_angle = userdata.rotation
      #rospy.loginfo("Rotating by {0}".format(change))
      rospy.loginfo("Rotating by {0}".format(userdata.rotation))
      #self.prev_angle = userdata.rotation
      goal = PoseArmGoal()
      goal.pose.header.frame_id = 'r_wrist_roll_link'
      goal.pose.pose.position.x = 0.0
      goal.pose.pose.position.y = 0.0
      goal.pose.pose.position.z = 0.0
      #temp = tt.quaternion_about_axis(change, (1,0,0)) # xaxis
      temp = tt.quaternion_about_axis(userdata.rotation, (1,0,0)) # xaxis
      #rospy.loginfo("Wrist quat: {0}".format(temp))
      goal.pose.pose.orientation.x = temp[0]
      goal.pose.pose.orientation.y = temp[1]
      goal.pose.pose.orientation.z = temp[2]
      goal.pose.pose.orientation.w = temp[3]

      self.right_arm_client.send_goal(goal)
      rospy.loginfo("Waiting for arms to clear")
      self.right_arm_client.wait_for_result()
      return 'succeeded'

class DialDisplay(State):
    def __init__(self):
        State.__init__(self, outcomes=['succeeded', 'failed', 'abort'], output_keys=['rotation'])

        self.layout = QVBoxLayout()
        self.dial = QDial()
        self.dial.valueChanged.connect(self.callback)
        self.dial.setNotchesVisible(True)
        self.dial.setWrapping(True)
        self.layout.addWidget(self.dial)

        self.button = QPushButton('done')
        self.button.clicked.connect(self.done)
        self.layout.addWidget(self.button)

        self.group = QGroupBox('knob control')
        self.group.setLayout(self.layout)

        self.pub = rospy.Publisher('knob/value', Float32)
        self.rotation = None

    def callback(self, value):
        self.pub.publish(float(value) / 99.0 * 3.14159 * 2.0)
        self.rotation = float(value) / 99.0 * 3.14159 * 2.0

    def done(self):
        self.group.hide()
        app.quit()

    def execute(self, userdata):
        rospy.loginfo("Dial display state")
        self.group.show()
        app.exec_()
        userdata.rotation = self.rotation
        return 'succeeded'

if __name__ == '__main__':
    rospy.init_node('finder')


    app = QApplication(sys.argv)
    machine = StateMachine(outcomes=['succeeded', 'failed', 'aborted'])
    with machine:
        StateMachine.add('CLEAR_ARM', ClearArm(), 
                                       transitions={'succeeded':'START', 
                                                    'abort':'aborted'})
        StateMachine.add('START', Click('Find', {'knobs':'succeeded', 'quit':'failed'}), 
                                       transitions={'succeeded':'SELECT_REGION',
                                                    'abort':'aborted'})
        StateMachine.add('SELECT_REGION', Display('/narrow_stereo/right/image_rect'), 
                                       transitions={'succeeded':'CALCULATE_PLANE', 
                                                    'abort':'aborted'})
        StateMachine.add('CALCULATE_PLANE', CalculatePlane('/narrow_stereo/points2', '/narrow_stereo/right/camera_info'), 
                                       transitions={'succeeded':'APPROACH_KNOB',
                                                    'abort':'aborted'})
        StateMachine.add('APPROACH_KNOB', PositionArm(), 
                                       transitions={'succeeded':'FINAL_APPROACH',
                                                    'abort':'aborted'})
        StateMachine.add('FINAL_APPROACH', FinalApproach(), 
                                       transitions={'succeeded':'GRIP_KNOB',
                                                    'abort':'aborted'})
        StateMachine.add('GRIP_KNOB', Gripper(True), 
                                       transitions={'succeeded':'DIAL',
                                                    'abort':'aborted'})
        StateMachine.add('DIAL', DialDisplay(), 
                                       transitions={'succeeded':'TURN_KNOB', 
                                                    'abort':'aborted'})
        StateMachine.add('TURN_KNOB', TurnKnob(), 
                                       transitions={'succeeded':'RELEASE_KNOB', 
                                                    'failed':'RELEASE_KNOB',
                                                    'abort':'aborted'})
        StateMachine.add('RELEASE_KNOB', Gripper(False), 
                                       transitions={'succeeded':'CLEAR_ARM',
                                                    'abort':'aborted'})


    try:
      rospy.loginfo("Starting state machine execution")
      machine.execute()
      rospy.spin()
    except KeyboardInterrupt:
      rospy.loginfo("Shutting down")
    cv.DestroyAllWindows()
